# -*- coding: utf-8 -*-
# @Time : 2021/11/17 15:14
# @Author : Xiang Qian Xiang Qian
# @Email : qianxljp@126.com
# @File : hrrp.py
# @Project : hrrp_jt
# !/usr/bin/env python

from jittor.dataset import Dataset
import pandas as pd
import numpy as np
from sklearn import preprocessing


class HRRP(Dataset):

    def __init__(self,
                 data_root=r'/datahub/hrrp_0.05/preprocessed/',
                 dset_name='train',
                 imb='Im0',
                 snr='non',
                 batch_size=64,
                 shuffle=True):
        # if you want to test resnet etc you should set input_channel = 3, because the net set 3 as the input dimensions
        super().__init__()
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.data_path = data_root
        self.snr = snr
        self.imb = imb
        self.d_name = dset_name
        self.data_id = '{}_{}'.format(self.imb, self.snr)

        path = self.data_path + self.data_id + '/'
        range_cells = pd.read_pickle(path + self.data_id + '_' + self.d_name + '.pkl')
        print(dset_name)
        print(range_cells['label'].value_counts())
        labels = range_cells.pop("label")
        range_cells = np.expand_dims(range_cells.values, 1)
        range_cells = np.expand_dims(range_cells, 2)
        # labels = preprocessing.LabelBinarizer().fit_transform(labels)
        # print(range_cells.shape)
        self.hrrp = {}
        self.hrrp["range_cells"], self.hrrp["labels"] = range_cells, labels.values

        assert (self.hrrp["range_cells"].shape[0] == self.hrrp["labels"].shape[0])
        self.total_len = self.hrrp["range_cells"].shape[0]
        # this function must be called
        self.set_attrs(batch_size=self.batch_size, total_len=self.total_len, shuffle=self.shuffle)

    def __getitem__(self, index):
        return self.hrrp["range_cells"][index], self.hrrp["labels"][index]

# a = HRRP()
# for batch_idx, (inputs, targets) in enumerate(a):
#     print(inputs.shape)
